# coding=utf-8 
'''
Created on 2018年2月21日
@author: heguofeng


clientA:                  webmaster                clientB:
start
get_nat_type(self)
register(period)                 
                          netusers(1m quit) 
get_peer()
   internet 
   fullcone
   restrictnat
   symmetricnat 
start_session                                      start_session
   send("100 1")
end_session
unregister
quit

connect    
   send(connectreq,userid,sharesecret)  receive(connectres,userid,sharesecret)
senddata
receivedata
disconnect
'''
import unittest

import random
import socket
import binascii
import logging
import struct


log = logging.getLogger("pystun")


def enable_logging():
    logging.basicConfig(filename="C:\\Users\\heguofeng\\log.txt")
    log.setLevel(logging.DEBUG)

stun_servers_list = (
                    "stun.stunprotocol.org",
                    "stun.sipgate.net",
                    "stun1.voiceeclipse.net",
                    #"stun.ekiga.net",
                    #'stunserver.org',
                    'stun.ideasip.com',
                    #'stun.softjoys.com',
                    'stun.voipbuster.com',
)

#stun attributes
MappedAddress = b'0001'
ResponseAddress = b'0002'
ChangeRequest = b'0003'
SourceAddress = b'0004'
ChangedAddress = b'0005'
Username = b'0006'
Password = b'0007'
MessageIntegrity = b'0008'
ErrorCode = b'0009'
UnknownAttribute = b'000A'
ReflectedFrom = b'000B'
XorOnly = b'0021'
XorMappedAddress = b'8020'
ServerName = b'8022'
SecondaryAddress = b'8050'  # Non standard extention

#types for a stun message
BindRequestMsg = '0001'
BindResponseMsg = '0101'
BindErrorResponseMsg = '0111'
SharedSecretRequestMsg = '0002'
SharedSecretResponseMsg = '0102'
SharedSecretErrorResponseMsg = '0112'

dictAttrToVal = {'MappedAddress': MappedAddress,
                 'ResponseAddress': ResponseAddress,
                 'ChangeRequest': ChangeRequest,
                 'SourceAddress': SourceAddress,
                 'ChangedAddress': ChangedAddress,
                 'Username': Username,
                 'Password': Password,
                 'MessageIntegrity': MessageIntegrity,
                 'ErrorCode': ErrorCode,
                 'UnknownAttribute': UnknownAttribute,
                 'ReflectedFrom': ReflectedFrom,
                 'XorOnly': XorOnly,
                 'XorMappedAddress': XorMappedAddress,
                 'ServerName': ServerName,
                 'SecondaryAddress': SecondaryAddress}

dictMsgTypeToVal = {
    'BindRequestMsg': BindRequestMsg,
    'BindResponseMsg': BindResponseMsg,
    'BindErrorResponseMsg': BindErrorResponseMsg,
    'SharedSecretRequestMsg': SharedSecretRequestMsg,
    'SharedSecretResponseMsg': SharedSecretResponseMsg,
    'SharedSecretErrorResponseMsg': SharedSecretErrorResponseMsg}

dictValToMsgType = {}

dictValToAttr = {}

Blocked = "Blocked"
OpenInternet = "Open Internet"
FullCone = "Full Cone"
SymmetricUDPFirewall = "Symmetric UDP Firewall"
RestrictNAT = "Restrict NAT"
RestrictPortNAT = "Restrict Port NAT"
SymmetricNAT = "Symmetric NAT"
ChangedAddressError = "Meet an error, when do Test1 on Changed IP and Port"



def bytes2addr(bytes):
    """Convert a hash to an address pair."""
    if len(bytes) != 8:
        raise ValueError("invalid bytes")
    host = socket.inet_ntoa(bytes[:4])
    port = struct.unpack("H", bytes[-4:-2])[0]  # unpack returns a tuple even if it contains exactly one item
    nat_type_id = struct.unpack("H", bytes[-2:])[0]
    target = (host, port)
    return target, nat_type_id

def _initialize():
    #items = dictAttrToVal.items()
    for k,v in dictAttrToVal.items():
        dictValToAttr.update({v:k})
    
    for k,v in dictMsgTypeToVal.items():
        dictValToMsgType.update({v: k})



class NATDetector(object):
    NAT_TYPE_LIST=["",
        OpenInternet,
        FullCone,
        RestrictNAT,
        RestrictPortNAT,
        SymmetricNAT,
        SymmetricUDPFirewall,
        Blocked,
        ChangedAddressError]  
    def __init__(self,sockfd,ip,port):
        self.sock = sockfd
        self.ip = ip
        self.port = port
        _initialize()
        pass
    
    @staticmethod
    def get_nat_type_id(nat_type_string):
        nat_type_id={    
                        OpenInternet :1,
                        FullCone :2,
                        RestrictNAT :3,
                        RestrictPortNAT :4,
                        SymmetricNAT :5,
                        SymmetricUDPFirewall :6,
                        Blocked:7,
                        ChangedAddressError :8,
                    }

        return nat_type_id[nat_type_string]

    @staticmethod
    def get_nat_type_string(nat_type_id):
        return NATDetector.NAT_TYPE_LIST[nat_type_id]
        

    def get_nat_info(self):
        #socket.setdefaulttimeout(2)
        nat_type,response = self._get_nat_type()
        print(nat_type,response)
        #socket.setdefaulttimeout(None)
        external_ip,external_port=None,None
        if response:
            external_ip = response.get("ExternalIP",None)
            external_port = response.get("ExternalPort",None)
        return nat_type,external_ip,external_port
        pass
    
    
    def _get_nat_type(self,server_ip=None,server_port=3478):
        if not server_ip:
            for server in stun_servers_list:
                nat_type,response = self._get_nat_type(server,3478)
                if nat_type != Blocked:
                    return nat_type,response
            return Blocked, None
        log.debug('Trying STUN host: %s' % server_ip)
        response = self.stun_request(server_ip,server_port,"")
        if not response['Resp']:
            return Blocked, None
        log.debug("Result: %s" % response)
        exIP = response['ExternalIP']
        exPort = response['ExternalPort']
        changedIP = response['ChangedIP']
        changedPort = response['ChangedPort']
        mayIPs = [response['ChangedIP'],response['SourceIP']]
        mayPorts = [response['ChangedPort'],response['SourcePort']]
        if response['ExternalIP'] == self.ip:
            changeRequest = ''.join([ChangeRequest.decode(), '0004', "00000006"])
            response = self.stun_request(server_ip,server_port,changeRequest,mayIPs,mayPorts)
            if response['Resp']:
                nat_type = OpenInternet
            else:
                nat_type = SymmetricUDPFirewall
        else:
            changeRequest = ''.join([ChangeRequest.decode(), '0004', "00000006"])
            log.debug("ChangeRequest")
            response = self.stun_request(server_ip,server_port,changeRequest,mayIPs,mayPorts)
            log.debug("Result: %s" % response)
            if response['Resp']:
                nat_type = FullCone
            else:
                log.debug("NoResponse,change IP and  port")
                response = self.stun_request(changedIP,changedPort)
                log.debug("Result: %s" % response)
                if not response['Resp']:
                    nat_type = ChangedAddressError
                else:
                    if exIP == response['ExternalIP'] and exPort == response['ExternalPort']:
                        changePortRequest = ''.join([ChangeRequest.decode(), '0004', "00000002"])
                        log.debug("ChangeRequest determine Restrict NAT or PortNAT")
                        response = self.stun_request(changedIP, server_port,changePortRequest,mayIPs,mayPorts)
                        log.debug("Result: %s" % response)
                        if response['Resp'] == True:
                            nat_type = RestrictNAT
                        else:
                            nat_type = RestrictPortNAT
                    else:
                        nat_type = SymmetricNAT

        return nat_type, response

    
    def _generate_tran_id(self):
        tran_id = ''
        for i in range(0,32):
            tran_id += random.choice('0123456789ABCDEF')  # RFC3489 128bits transaction ID
        return tran_id
    
    def _get_ip_port(self,buf,base):
        port = int(binascii.b2a_hex(buf[base + 6:base + 8]), 16)
        ip = ".".join([ str(int(binascii.b2a_hex(buf[base + 8:base + 9]), 16)),
                        str(int(binascii.b2a_hex(buf[base + 9:base + 10]), 16)),
                        str(int(binascii.b2a_hex(buf[base + 10:base + 11]), 16)),
                        str(int(binascii.b2a_hex(buf[base + 11:base + 12]), 16))])
        return ip,port

    def stun_request(self, host, port, send_data="",mayhosts=[],mayports=[]):
        retVal = {'Resp': False, 'ExternalIP': None, 'ExternalPort': None,
                  'SourceIP': None, 'SourcePort': None, 'ChangedIP': None,
                  'ChangedPort': None}
        str_len = "%#04d" % (len(send_data) / 2)
        tranid = self._generate_tran_id()
        str_data = ''.join([BindRequestMsg, str_len, tranid, send_data])
        data = binascii.a2b_hex(str_data)
        recvCorr = False
        while not recvCorr:
            recieved = False
            count = 3
            while not recieved:
                log.debug("sendto %s" % str((host, port)))
                try:
                    ip = socket.gethostbyname(host)
                    self.sock.sendto(data, (host, port))
                except socket.gaierror:
                    retVal['Resp'] = False
                    return retVal
                try:
                    buf, addr = self.sock.recvfrom(2048)
                    log.debug("recvfrom: %s" % str(addr))
                    print(addr,mayhosts,mayports)
                    if mayhosts:
                        if addr[0] in mayhosts:
                            recieved = True
                    if mayports:
                        if addr[1] in mayports:
                            recieved = True
                    if not mayhosts and not mayports and ip == addr[0]:
                        recieved = True
                except Exception:
                    recieved = False
                    if count > 0:
                        count -= 1
                    else:
                        retVal['Resp'] = False
                        return retVal
            msgtype = binascii.b2a_hex(buf[0:2]).decode()
            bind_resp_msg = dictValToMsgType[msgtype] == "BindResponseMsg"
            tranid_match = tranid.upper() == binascii.b2a_hex(buf[4:20]).upper().decode()
            log.debug("%s %s %s %s", msgtype,bind_resp_msg,binascii.b2a_hex(buf[4:20]).upper().decode(),tranid.upper())
            if bind_resp_msg and tranid_match:
                recvCorr = True
                retVal['Resp'] = True
                len_message = int(binascii.b2a_hex(buf[2:4]), 16)
                len_remain = len_message
                base = 20
                while len_remain:
                    attr_type = binascii.b2a_hex(buf[base:(base + 2)])
                    attr_len = int(binascii.b2a_hex(buf[(base + 2):(base + 4)]),16)

                    if attr_type == MappedAddress:  # first two bytes: 0x0001
                        ip,port = self._get_ip_port(buf,base)
                        retVal['ExternalIP'] = ip
                        retVal['ExternalPort'] = port
                    if attr_type == SourceAddress:
                        ip,port = self._get_ip_port(buf,base)
                        retVal['SourceIP'] = ip
                        retVal['SourcePort'] = port
                    if attr_type == ChangedAddress:
                        ip,port = self._get_ip_port(buf,base)
                        retVal['ChangedIP'] = ip
                        retVal['ChangedPort'] = port
                    #if attr_type == ServerName:
                        #serverName = buf[(base+4):(base+4+attr_len)]
                    log.debug("%s",attr_type.decode())
                    base = base + 4 + attr_len
                    len_remain = len_remain - (4 + attr_len)
            else:
                pass
              
        return retVal

  
    


class Test(unittest.TestCase):


    def test_get_nat_info(self):
        self.ip,self.port = "0.0.0.0",54320
        self.sockfd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.sockfd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        logging.info("Socket bind to %s,%d"%(self.ip,self.port))
        self.sockfd.bind((self.ip,self.port))
        natdetector = NATDetector(self.sockfd,self.ip,self.port)
        nat_type,self.ip,self.port = natdetector.get_nat_info()
        print(nat_type,self.ip,self.port)

        pass

if __name__ == "__main__":
    ip,port = "0.0.0.0",54320
    socket.setdefaulttimeout(2)
    sockfd = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    sockfd.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    print("Socket bind to %s,%d"%(ip,port))
    sockfd.bind((ip,port))
    natdetector = NATDetector(sockfd,ip,port)
    nat_type,ip,port = natdetector.get_nat_info()
    print(nat_type,ip,port)
    socket.setdefaulttimeout(None)
